import os
import fileinput
from tqdm import tqdm
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import pandas as pd
import numpy as np



def generate_training_label():
    label_pd = pd.read_csv("label_pd.csv")

    label_pd.head()
    type_list = label_pd["type"].unique().tolist()

    type_list.remove("DontCare")
    type_list = ["DontCare"] + type_list
    type_dict = {t: i for i, t in enumerate(type_list)}
    label_pd["type_id"] = label_pd.apply(
        lambda row: type_dict[row["type"]], axis=1
    )
    train_label = label_pd.loc[
        :,
        [
            "#id",
            "type_id",
            "bbox_left",
            "bbox_top",
            "bbox_right",
            "bbox_bottom",
        ],
    ]
    neg_item = pd.DataFrame([0, 0, 0, 0, 0, 0]).T
    neg_item.columns = train_label.columns
    train_label = pd.concat(
        [neg_item, train_label], ignore_index=True
    )
    train_label.head()

    training_label = []
    for i in tqdm(range(train_label["#id"].max())):
        tmp = train_label[train_label["#id"] == i].iloc[:6]
        img_size = Image.open(get_pic_name(43)).size
        bbox_fixer = np.array([1] + list(img_size) * 2)
        training_label.append(
            pd.concat(
                [tmp] + [neg_item] * (6 - len(tmp)), ignore_index=True
            )
            .iloc[:, 1:]
            .to_numpy()
            / bbox_fixer
        )
    training_label = torch.tensor(training_label)
    torch.save(training_label, "training_label.pt")


def label_recompose(label_dir="training/label_2/"):
    Columns_name = [
        "type",
        "truncated",
        "occluded",
        "alpha",
        "bbox_lift",
        "bbox_top",
        "bbox_right",
        "bbox_bottom",
        "height",
        "width",
        "length",
        "pos_x",
        "pos_y",
        "pos_z",
        "rot_y",
    ]
    label_files = os.listdir(label_dir)
    fileinput.close()
    label_df_ones = []
    for file in tqdm(label_files):
        for line in fileinput.input(os.path.join(label_dir, file)):
            item = dict(zip(Columns_name, line.strip().split(" ")))
            item["#id"] = int(
                os.path.basename(fileinput.filename()).split(".")[0]
            )
            label_df_ones.append(
                pd.DataFrame([item], columns=["#id"] + Columns_name)
            )
        #             print(item)
        fileinput.close()
    label_df = pd.concat(label_df_ones)
    label_df.to_csv("label_df.csv", index=False)


def calib_recompose():
    column_items = {
        "P0": 12,
        "P1": 12,
        "P2": 12,
        "P3": 12,
        "R0_rect": 9,
        "Tr_velo_to_cam": 12,
        "Tr_imu_to_velo": 12,
    }
    calib_dir = "data_object_calib/training/calib/"
    calib_dfs = []

    for k, file in tqdm(enumerate(os.listdir(calib_dir))):
        print(k)
        with open(os.path.join(calib_dir, file)) as f:
            one_file = f.readlines()
            item = {}
            for i, key in enumerate(column_items.keys()):
                values = [k] + one_file[i].split(":")[
                    1
                ].strip().split(" ")
                keys = ["#id"] + [
                    key + "_" + str(k)
                    for k in range(column_items[key])
                ]
                item.update(dict(zip(keys, values)))
            calib_dfs.append(pd.DataFrame([item]))

    calib_df = pd.concat(calib_dfs, ignore_index=True)
    calib_df.to_csv("calib_df.csv")


def compute_3D_box_cam2(h, w, l, x, y, z, yaw):
    """
    Return:3Xn in cam2 coordinate
    """
    # 建立旋转矩阵R
    R = np.array(
        [
            [np.cos(yaw), 0, np.sin(yaw)],
            [0, 1, 0],
            [-np.sin(yaw), 0, np.cos(yaw)],
        ]
    )
    # print(R)
    # 计算8个顶点坐标
    x_corners = [
        l / 2,
        l / 2,
        -l / 2,
        -l / 2,
        l / 2,
        l / 2,
        -l / 2,
        -l / 2,
        0,
    ]
    y_corners = [0, 0, 0, 0, -h, -h, -h, -h, -h / 2]
    z_corners = [
        w / 2,
        -w / 2,
        -w / 2,
        w / 2,
        w / 2,
        -w / 2,
        -w / 2,
        w / 2,
        0,
    ]
    # 使用旋转矩阵变换坐标
    corners_3d_cam2 = np.dot(
        R, np.vstack([x_corners, y_corners, z_corners])
    )
    # 最后在加上中心点
    corners_3d_cam2 += np.vstack([x, y, z])
    return corners_3d_cam2


def draw_box(ax, vertices, axes=[0, 1, 2], color="red"):
    """
    Draws a bounding 3D box in a pyplot axis.

    Parameters
    ----------
    pyplot_axis : Pyplot axis to draw in.
    vertices    : Array 8 box vertices containing x, y, z coordinates.
    axes        : Axes to use. Defaults to `[0, 1, 2]`, e.g. x, y and z axes.
    color       : Drawing color. Defaults to `black`.
    """
    vertices = vertices[axes, :]
    connections = [
        [0, 1],
        [1, 2],
        [2, 3],
        [3, 0],  # Lower plane parallel to Z=0 plane
        [4, 5],
        [5, 6],
        [6, 7],
        [7, 4],  # Upper plane parallel to Z=0 plane
        [0, 4],
        [1, 5],
        [2, 6],
        [3, 7],  # Connections between upper and lower planes
    ]
    for connection in connections:
        ax.plot(*vertices[:, connection], c=color, lw=2)
    print(vertices[:, connections])

    ax.scatter(*vertices[:, -1], s=100, c="y")

def show_pic(label_df, calib_df, img_idx=32):

    fig = plt.figure(figsize=(12, 5))
    ax = fig.add_subplot(111)
    im = (
        "data_object_image_2/training/image_2/"
        + "0" * (6 - len(str(img_idx)))
        + str(img_idx)
        + ".png"
    )
    img = mpimg.imread(im)

    label_3d_one_pic = label_df[
        (label_df["#id"] == img_idx) & (label_df["occluded"] != -1)
    ].loc[
        :,
        [
            "height",
            "width",
            "length",
            "pos_x",
            "pos_y",
            "pos_z",
            "rot_y",
        ],
    ]
    calib_one = calib_df[calib_df["#id"] == img_idx]
    for i in range(len(label_3d_one_pic)):
        corners_3d_cam2 = compute_3D_box_cam2(
            *label_3d_one_pic.iloc[i].values
        )
        # print(corners_3d_cam2)
        screen_3d = np.dot(
            calib_one.loc[:, calib_one.columns.str.startswith("P2")]
            .to_numpy()
            .reshape(3, 4)[:, :-1],
            corners_3d_cam2,
        )

        screen_3d /= screen_3d[2]
        # print(screen_3d)
        draw_box(ax, screen_3d, axes=[0, 1])

    ax.imshow(img)


def show_pic_info(img_idx, label_df, calib_df):
    print(label_df[label_df["#id"] == img_idx])
    print(calib_df[calib_df["#id"] == img_idx].to_dict())
